﻿#include "Aes.h"
#include "XyNet/Common/CrossObjectHolder.h"
#include "XyNet/Common/Exception.h"
#include "XyNet/Crypto/Sha.h"
#include "XyNet/Crypto/X917Random.h"

using namespace CryptoPP;

namespace XyNet
{
	Aes::KeyType Aes::_defaultKeyType = Aes128;
	const int Aes::_ivByteLength = 16;

	Aes::Aes(KeyType keyType)
	{
		setKeyType(keyType);
	}

	Aes::Aes(const Bytes& key, const Bytes& iv)
	{
		setKeyAndIv(key, iv);
	}

	Aes::KeyType Aes::getKeyType() const
	{
		assert(_keyType != AesDefault);
		return _keyType;
	}

	void Aes::setKeyType(KeyType keyType)
	{
		if (keyType == AesDefault)
			_keyType = _defaultKeyType;
		else
			_keyType = keyType;

		_key.resize(getKeyByteLength());
		X917Random::getDefault()->GenerateBlock((byte*)_key.data(), _key.size());
		_iv.resize(_ivByteLength);
		X917Random::getDefault()->GenerateBlock((byte*)_iv.data(), _iv.size());
		applyKeyAndIv();
	}

	const Bytes& Aes::getKey() const
	{
		return _key;
	}

	int Aes::getKeyByteLength() const
	{
		return getKeyType() / 8;
	}

	int Aes::getKeyBitLength() const
	{
		return getKeyType();
	}

	const Bytes& Aes::getIv() const
	{
		return _iv;
	}

	int Aes::getIvByteLength() const
	{
		return _ivByteLength;
	}

	int Aes::getIvBitLength() const
	{
		return _ivByteLength * 8;
	}

	void Aes::setKeyAndIv(const Bytes& key, const Bytes& iv)
	{
		_key = key;
		_iv = iv;
		if (_iv.empty())
			_iv.assign(_ivByteLength, 0);
		applyKeyAndIv();
	}

	void Aes::setIv(const Bytes& iv)
	{
		_iv = iv;
		applyKeyAndIv();
	}

	void Aes::setRandomIv()
	{
		X917Random::getDefault()->GenerateBlock((byte*)_iv.data(), _iv.size());
		applyKeyAndIv();
	}

	Bytes Aes::encrypt(const Bytes& input)
	{
		Bytes output(input.size(), 0);
		auto encryptor = _encryptor;
		encryptor.ProcessData((byte*)output.data(), (byte*)input.data(), input.size());
		return output;
	}

	Bytes Aes::decrypt(const Bytes& input)
	{
		Bytes output(input.size(), 0);
		auto decryptor = _decryptor;
		decryptor.ProcessData((byte*)output.data(), (byte*)input.data(), input.size());
		return output;
	}

	Aes::KeyType Aes::getDefaultKeyType()
	{
		return _defaultKeyType;
	}

	void Aes::setDefaultKeyType(KeyType keyType)
	{
		assert(keyType != AesDefault);
		_defaultKeyType = keyType;
	}

	void Aes::applyKeyAndIv()
	{
		Sha sha(Sha::DigestType(getKeyBitLength() * 2));
		Bytes digestKey = sha.calculateDigest(_key).substr(getKeyByteLength());
		sha.setDigestType(Sha::DigestType(getIvBitLength() * 2));
		Bytes digestIv = sha.calculateDigest(_iv).substr(getIvByteLength());

		_encryptor.SetKeyWithIV((byte*)digestKey.data(), digestKey.size(),
			(byte*)digestIv.data(), digestIv.size());
		_decryptor.SetKeyWithIV((byte*)digestKey.data(), digestKey.size(),
			(byte*)digestIv.data(), digestIv.size());
	}
}

//////////////////////////////////////////////////////////////////////////

namespace XyNet
{
	XYNET_IMPLEMENT_BASE_SINGLETON(CrossObjectHolder<Aes>);

	XYNET_API Aes* AesCreateWithKeyType(Aes::KeyType keyType)
	{
		try
		{
			return CrossObjectHolder<Aes>::getInstance()->add(make_shared<Aes>(keyType));
		}
		catch (const exception& e)
		{
			HandleException(e);
			return nullptr;
		}
	}

	XYNET_API Aes* AesCreateWithKeyAndIv(Bytes* key, Bytes* iv)
	{
		try
		{
			return CrossObjectHolder<Aes>::getInstance()->add(make_shared<Aes>(*key, *iv));
		}
		catch (const exception& e)
		{
			HandleException(e);
			return nullptr;
		}
	}

	XYNET_API void AesDestroy(Aes* aes)
	{
		try
		{
			CrossObjectHolder<Aes>::getInstance()->remove(aes);
		}
		catch (const exception& e)
		{
			HandleException(e);
		}
	}

	XYNET_API Aes::KeyType AesGetKeyType(Aes* aes)
	{
		try
		{
			return aes->getKeyType();
		}
		catch (const exception& e)
		{
			HandleException(e);
			return Aes::KeyType(0);
		}
	}

	XYNET_API void AesSetKeyType(Aes* aes, Aes::KeyType keyType)
	{
		try
		{
			aes->setKeyType(keyType);
		}
		catch (const exception& e)
		{
			HandleException(e);
		}
	}

	XYNET_API Bytes* AesGetKey(Aes* aes)
	{
		try
		{
			return BytesHold(aes->getKey());
		}
		catch (const exception& e)
		{
			HandleException(e);
			return nullptr;
		}
	}

	XYNET_API Bytes* AesGetIv(Aes* aes)
	{
		try
		{
			return BytesHold(aes->getIv());
		}
		catch (const exception& e)
		{
			HandleException(e);
			return nullptr;
		}
	}

	XYNET_API void AesSetKeyAndIv(Aes* aes, Bytes* key, Bytes* iv)
	{
		try
		{
			aes->setKeyAndIv(*key, *iv);
		}
		catch (const exception& e)
		{
			HandleException(e);
		}
	}

	XYNET_API Bytes* AesEncrypt(Aes* aes, Bytes* input)
	{
		try
		{
			return BytesHold(aes->encrypt(*input));
		}
		catch (const exception& e)
		{
			HandleException(e);
			return nullptr;
		}
	}

	XYNET_API Bytes* AesDecrypt(Aes* aes, Bytes* input)
	{
		try
		{
			return BytesHold(aes->decrypt(*input));
		}
		catch (const exception& e)
		{
			HandleException(e);
			return nullptr;
		}
	}

	XYNET_API Aes::KeyType AesGetDefaultKeyType()
	{
		try
		{
			return Aes::getDefaultKeyType();
		}
		catch (const exception& e)
		{
			HandleException(e);
			return Aes::KeyType(0);
		}
	}

	XYNET_API void AesSetDefaultKeyType(Aes::KeyType keyType)
	{
		try
		{
			return Aes::setDefaultKeyType(keyType);
		}
		catch (const exception& e)
		{
			HandleException(e);
		}
	}
}